{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J1w-zrxAUr4l",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "# Quick Start Tutorial of the basic functionality of MetaDrive\n",
        "\n",
        "Welcome to MetaDrive!\n",
        "\n",
        "MetaDrive v0.2.5 supports two running modes:\n",
        "\n",
        "1. **With 3D rendering functionality**: MetaDrive can easily install and run on personal computer, but may need special treatments for 3D rendering in headless machines and cloud servers.\n",
        "\n",
        "2. **Without 3D rendering functionality**: MetaDrive can easily be installed and run on any machine. In this Colab notebook, we mainly demonstrate MetaDrive in this mode and the renderer will be the **2D** **Pygame** renderer.\n",
        "\n",
        "In this tutorial, we will navigate you through the installation and some basic functionality of MetaDrive!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2efvTXdHVptN",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "## Installation\n",
        "\n",
        "You can install MetaDrive easily."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "65J2iQKpUQ1B",
        "pycharm": {
          "name": "#%%\n"
        },
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "#@title Collect the MetaDrive\n",
        "\n",
        "%pip install git+https://github.com/metadriverse/metadrive.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XGXIZt4TsPhT"
      },
      "outputs": [],
      "source": [
        "# @title Test Installation\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "from metadrive import MetaDriveEnv\n",
        "from metadrive.examples import expert\n",
        "\n",
        "env = MetaDriveEnv()\n",
        "print(\"\\nThe action space: {}\".format(env.action_space))\n",
        "print(\"\\nThe observation space: {}\\n\".format(env.observation_space))\n",
        "try:\n",
        "    env.reset()\n",
        "    for i in range(100):\n",
        "        env.step(env.action_space.sample())\n",
        "    print(\"Successfully run MetaDrive\")\n",
        "except:\n",
        "    raise ValueError(\"Fail to run MetaDrive!\")\n",
        "finally:\n",
        "    env.close()\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QhteKnQ5kqEg",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "## Basic Functionality"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "898SwWaHkqEg",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "# @title A minimalist example of using MetaDrive with the built-in PPO expert policy\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "import os\n",
        "os.environ['SDL_VIDEODRIVER']='dummy'\n",
        "# Note: this step is only for cheating Colab\n",
        "\n",
        "from metadrive import MetaDriveEnv\n",
        "from metadrive.examples import expert\n",
        "\n",
        "# Set the envrionment config\n",
        "config = {\"start_seed\": 1000,\n",
        "          \"num_scenarios\":1,\n",
        "          \"traffic_density\":0.1,\n",
        "          }\n",
        "\n",
        "env = MetaDriveEnv(config)\n",
        "\n",
        "print(\"Starting the environment ...\\n\")\n",
        "\n",
        "ep_reward = 0.0\n",
        "obs, info = env.reset()\n",
        "for i in range(1000):\n",
        "    obs, reward, terminated, truncated, info = env.step(expert(env.agent))\n",
        "    ep_reward += reward\n",
        "    env.render(mode=\"top_down\", screen_record=True, screen_size=(500, 500))\n",
        "    if terminated or truncated:\n",
        "        print(\"Arriving Destination: {}\".format(info[\"arrive_dest\"]))\n",
        "        print(\"\\nEpisode reward: \", ep_reward)\n",
        "        break\n",
        "\n",
        "print(\"\\nThe last returned information: {}\".format(info))\n",
        "env.top_down_renderer.generate_gif()\n",
        "env.close()\n",
        "print(\"\\nMetaDrive successfully run!\")\n",
        "\n",
        "from IPython.display import Image\n",
        "Image(open(\"demo.gif\", 'rb').read())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I2uN9XB4sPhW"
      },
      "outputs": [],
      "source": [
        "# @title PPO expert policy can also drive in the safety-critical SafeMetaDrive environment\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "import os\n",
        "os.environ['SDL_VIDEODRIVER']='dummy'\n",
        "# Note: this step is only for cheating Colab\n",
        "\n",
        "from metadrive import SafeMetaDriveEnv\n",
        "from metadrive.examples import expert\n",
        "\n",
        "# Set the environment config\n",
        "config = {\"start_seed\": 0,\n",
        "          \"num_scenarios\":1,\n",
        "          \"accident_prob\":1.0,\n",
        "          \"traffic_density\":0.25,\n",
        "          }\n",
        "\n",
        "env = SafeMetaDriveEnv(config)\n",
        "\n",
        "print(\"Starting the environment ...\\n\")\n",
        "\n",
        "ep_reward = 0.0\n",
        "ep_cost = 0.0\n",
        "obs, info = env.reset()\n",
        "frames = []\n",
        "for i in range(1000):\n",
        "    obs, reward, terminated, truncated, info = env.step(expert(env.agent))\n",
        "    ep_reward += reward\n",
        "    ep_cost += info[\"cost\"]\n",
        "    frame = env.render(mode=\"top_down\", window=False, screen_size=(500, 500))\n",
        "    frames.append(frame)\n",
        "    if terminated or truncated:\n",
        "        print(\"Arriving Destination: {}\".format(info[\"arrive_dest\"]))\n",
        "        print(\"\\nEpisode reward: \", ep_reward)\n",
        "        print(\"\\nEpisode cost: \", ep_cost)\n",
        "        break\n",
        "\n",
        "env.close()\n",
        "print(\"\\nMetaDrive successfully run!\")\n",
        "\n",
        "# render image\n",
        "print(\"\\nGenerate gif...\")\n",
        "import pygame\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "imgs = [frame for frame in frames]\n",
        "imgs = [Image.fromarray(img) for img in imgs]\n",
        "imgs[0].save(\"demo.gif\", save_all=True, append_images=imgs[1:], duration=50, loop=0)\n",
        "print(\"\\nOpen gif...\")\n",
        "from IPython.display import Image\n",
        "Image(open(\"demo.gif\", 'rb').read())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J3VKSy33sPhX"
      },
      "outputs": [],
      "source": [
        "# @title Multi-agent Environment Visualization\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "import os\n",
        "os.environ['SDL_VIDEODRIVER']='dummy'\n",
        "# Note: this step is only for cheating Colab\n",
        "\n",
        "from metadrive.envs.marl_envs import MultiAgentRoundaboutEnv, MultiAgentBottleneckEnv, MultiAgentIntersectionEnv, MultiAgentParkingLotEnv, MultiAgentTollgateEnv\n",
        "from metadrive.examples import expert\n",
        "\n",
        "env_classes = [MultiAgentRoundaboutEnv,\n",
        "        MultiAgentBottleneckEnv,\n",
        "        MultiAgentIntersectionEnv,\n",
        "        MultiAgentParkingLotEnv,\n",
        "        MultiAgentTollgateEnv]\n",
        "\n",
        "frames = []\n",
        "for env_class in env_classes:\n",
        "    env = env_class()\n",
        "    print(\"Starting the environment {}\\n\".format(env))\n",
        "    env.reset()\n",
        "    tm={\"__all__\":False}\n",
        "    for i in range(100):\n",
        "        if tm[\"__all__\"]:\n",
        "            frames.append(frame)\n",
        "            continue\n",
        "        action = env.action_space.sample()\n",
        "        for a in action.values():\n",
        "            a[-1] = 1.0\n",
        "        o,r,tm,tc,i = env.step(action)\n",
        "        frame = env.render(mode=\"top_down\",\n",
        "                           scaling=4, # 4 pixels per meter\n",
        "                           camera_position=env.current_map.get_center_point(),\n",
        "                           screen_size=(500, 500))\n",
        "        frames.append(frame)\n",
        "    env.close()\n",
        "\n",
        "# render image\n",
        "print(\"\\nGenerate gif...\")\n",
        "import pygame\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "imgs = [frame for frame in frames]\n",
        "imgs = [Image.fromarray(img) for img in imgs]\n",
        "imgs[0].save(\"demo.gif\", save_all=True, append_images=imgs[1:], duration=50, loop=0)\n",
        "print(\"\\nOpen gif...\")\n",
        "from IPython.display import Image\n",
        "Image(open(\"demo.gif\", 'rb').read())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PPkiDVgrsPhY"
      },
      "outputs": [],
      "source": [
        "# @title Real-world Scenario Environment Visualization\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "import argparse\n",
        "import random\n",
        "from metadrive.policy.replay_policy import ReplayEgoCarPolicy\n",
        "from metadrive.constants import HELP_MESSAGE\n",
        "from metadrive.engine.asset_loader import AssetLoader\n",
        "from metadrive.envs import ScenarioEnv\n",
        "\n",
        "\n",
        "class DemoEnv(ScenarioEnv):\n",
        "    def reset(self, seed=None):\n",
        "        if self.engine is not None:\n",
        "            seeds = [i for i in range(self.config[\"num_scenarios\"])]\n",
        "            seeds.remove(self.current_seed)\n",
        "            seed = random.choice(seeds)\n",
        "        return super(DemoEnv, self).reset(seed=seed)\n",
        "\n",
        "\n",
        "extra_args = dict(film_size=(1200, 1200))\n",
        "asset_path = AssetLoader.asset_path\n",
        "\n",
        "try:\n",
        "    env = DemoEnv(\n",
        "        {\n",
        "            \"manual_control\": False,\n",
        "            \"reactive_traffic\": False,\n",
        "            \"use_render\": False,\n",
        "            \"agent_policy\": ReplayEgoCarPolicy,\n",
        "            \"data_directory\": AssetLoader.file_path(asset_path, \"waymo\", unix_style=False),\n",
        "            \"num_scenarios\": 3\n",
        "        }\n",
        "    )\n",
        "    o, _ = env.reset()\n",
        "    frames = []\n",
        "    for i in range(1, 100000):\n",
        "        o, r, tm, tc, info = env.step([1.0, 0.])\n",
        "        frame=env.render(\n",
        "            mode=\"top_down\",\n",
        "            **extra_args\n",
        "        )\n",
        "        frames.append(frame)\n",
        "        if tm or tc:\n",
        "            env.reset()\n",
        "            break\n",
        "\n",
        "except Exception as e:\n",
        "    raise e\n",
        "finally:\n",
        "    env.close()\n",
        "\n",
        "# render image\n",
        "print(\"\\nGenerate gif...\")\n",
        "import pygame\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "imgs = [frame for frame in frames]\n",
        "imgs = [Image.fromarray(img) for img in imgs]\n",
        "imgs[0].save(\"demo.gif\", save_all=True, append_images=imgs[1:], duration=50, loop=0)\n",
        "print(\"\\nOpen gif...\")\n",
        "from IPython.display import Image\n",
        "Image(open(\"demo.gif\", 'rb').read())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ClZzNVQUkqEh",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "## Map Generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NUawlJubkqEh",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "# @title Draw the generated maps in top-down view\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "\n",
        "import random\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from metadrive import MetaDriveEnv\n",
        "from metadrive.utils.draw_top_down_map import draw_top_down_map\n",
        "\n",
        "env = MetaDriveEnv(config=dict(\n",
        "    num_scenarios=100,\n",
        "    map=7,\n",
        "    start_seed=random.randint(0, 1000)\n",
        "))\n",
        "\n",
        "fig, axs = plt.subplots(4, 4, figsize=(10, 10), dpi=100)\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        env.reset()\n",
        "        m = draw_top_down_map(env.current_map)\n",
        "        ax = axs[i][j]\n",
        "        ax.imshow(m, cmap=\"bone\")\n",
        "        ax.set_xticks([])\n",
        "        ax.set_yticks([])\n",
        "fig.suptitle(\"Bird's-eye view of generated maps\")\n",
        "plt.show()\n",
        "\n",
        "env.close()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HyL36-nekqEh",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "# @title Draw the generated maps in top-down view with fixed block sequence\n",
        "# @markdown You can also specify the road block sequence then randomize the block parameters.\n",
        "# @markdown Please refer to [documentation](https://metadrive-simulator.readthedocs.io/en/latest/env_config.html#map-config) for the meaning of the map string.\n",
        "from metadrive.engine.engine_utils import close_engine\n",
        "close_engine()\n",
        "# NOTE: usually you don't need the above lines. It is only for avoiding a potential bug when running on colab\n",
        "\n",
        "\n",
        "import random\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from metadrive import MetaDriveEnv\n",
        "\n",
        "env = MetaDriveEnv(config=dict(\n",
        "    num_scenarios=100,\n",
        "    map=\"CrTRXOS\",\n",
        "    start_seed=random.randint(0, 1000)\n",
        "))\n",
        "\n",
        "fig, axs = plt.subplots(4, 4, figsize=(10, 10), dpi=100)\n",
        "for i in range(4):\n",
        "    for j in range(4):\n",
        "        env.reset()\n",
        "        m = draw_top_down_map(env.current_map)\n",
        "        ax = axs[i][j]\n",
        "        ax = axs[i][j]\n",
        "        ax.imshow(m, cmap=\"bone\")\n",
        "        ax.set_xticks([])\n",
        "        ax.set_yticks([])\n",
        "fig.suptitle(\"Bird's-eye view of generated maps\")\n",
        "plt.show()\n",
        "\n",
        "env.close()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5OoFwaKsPhb"
      },
      "source": [
        "## 3D Renderer\n",
        "For using the 3D renderer, just set `use_render=True` in the environment config. Or you can turn on the `image_observation`, which will make an offscreen renderer so we can access the rendered result by capturing frames."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jTuSZ7vKsPhb"
      },
      "outputs": [],
      "source": [
        "from metadrive.envs.metadrive_env import MetaDriveEnv\n",
        "import os\n",
        "\n",
        "cfg=dict(image_observation=True,\n",
        "         show_terrain=not os.getenv('TEST_IPYNB'),\n",
        "         sensors={\"main_camera\": ()},\n",
        "         vehicle_config=dict(image_source=\"main_camera\"),\n",
        "         window_size=(84, 60))\n",
        "\n",
        "env=MetaDriveEnv(cfg)\n",
        "try:\n",
        "    env.reset()\n",
        "    for _ in range(1 if os.getenv('TEST_DOC') else 5):\n",
        "        env.step([0,0])\n",
        "    env.capture(\"3d_renderer.png\")\n",
        "\n",
        "finally:\n",
        "    env.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cSHp8WGmsPhc"
      },
      "outputs": [],
      "source": [
        "from IPython.display import Image\n",
        "Image(open(\"3d_renderer.png\", \"rb\").read())"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}